/******************************************************************************
	This file is a part of the Snapshot Bioinformatics Project

	Copyright (C) 2012-2013 Jarom Schow

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>.
******************************************************************************/

#include "GenBank.h"

#include <QtCore/QCoreApplication>
#include <QtCore/QBuffer>
#include <QtCore/QUrl>

#include "qftp.h"

#include "Core.h"

#define QFTP

using namespace Bio;

const QString	c_ftpHost = "ftp.ncbi.nih.gov";
const QString	c_ftpDir = "genbank";
const QString	c_releaseNumberFile = "GB_Release_Number";

struct DivisionDefs {
	DivisionDefs( const char *abbr, const char *desc, const char *prefix ) : m_abbr( abbr ), m_desc( desc ), m_filePrefix( prefix )
	{}

	QString	m_abbr, m_desc, m_filePrefix;
};

const DivisionDefs	c_defs[] = {
	DivisionDefs( "PRI", "Primate sequences", "gbpri" ),
	DivisionDefs( "ROD", "Rodent sequences", "gbrod" ),
	DivisionDefs( "MAM", "Other Mammalian sequences", "gbmam" ),
	DivisionDefs( "VRT", "Other Vertebrate sequences", "gbvrt" ),
	DivisionDefs( "INV", "Invertebrate sequences", "gbinv" ),
	DivisionDefs( "PLN", "Plant, Fungal, and Algal sequences", "gbpln" ),
	DivisionDefs( "BCT", "Bacterial sequences", "gbbct" ),
	DivisionDefs( "VRL", "Viral sequences", "gbvrl" ),
	DivisionDefs( "PHG", "Bacteriophage sequences", "gbphg" ),
	DivisionDefs( "SYN", "Synthetic sequences", "gbsyn" ),
	DivisionDefs( "UNA", "Unannotated sequences", "gbuna" ),
	DivisionDefs( "PAT", "Patent sequences", "gbpat" ),
	DivisionDefs( "ENV", "Environmental sampling sequences", "gbenv" ),
	DivisionDefs( "CON", "Constructed sequences", "gbcon" ),
	DivisionDefs( "EST", "Expressed sequence tags", "gbest" ),
	DivisionDefs( "STS", "Sequence tagged sites", "gbsts" ),
	DivisionDefs( "GSS", "Genome survey sequences", "gbgss" ),
	DivisionDefs( "HTG", "High throughput genomic sequences", "gbhtg" ),
	DivisionDefs( "HTC", "High throughput cDNA sequences", "gbhtc" )
};

///////////////////////////////////////////////////////////////////////////////
//	GenBank
///////////////////////////////////////////////////////////////////////////////

/**
	@class GenBank
	@brief A utility class for accessing the GenBank ftp.

	Provides functions for accessing the files on the GenBank ftp server.
**/

/**
	@enum GenBank::Division
	The enumerated divisions of GenBank.
**/

/**
	@var GenBank::Division GenBank::Primate
	The PRI division of GenBank containing Primate sequence data.
**/

/**
	@var GenBank::Division GenBank::Rodent
	The ROD division of GenBank containing Rodent sequence data.
**/

/**
	@var GenBank::Division GenBank::Mammal
	The MAM division of GenBank containing Mammal sequence data.
**/

/**
	@var GenBank::Division GenBank::Vertebrate
	The VRT division of GenBank containing Vertebrate sequence data.
**/

/**
	@var GenBank::Division GenBank::Invertebrate
	The INV division of GenBank containing Invertebrate sequence data.
**/

/**
	@var GenBank::Division GenBank::Plant
	The PLN division of GenBank containing Plant sequence data.
**/

/**
	@var GenBank::Division GenBank::Bacterial
	The BAC division of GenBank containing Bacterial sequence data.
**/

/**
	@var GenBank::Division GenBank::Viral
	The VRL division of GenBank containing Viral sequence data.
**/

/**
	@var GenBank::Division GenBank::Phage
	The PHG division of GenBank containing Bacteriophage sequence data.
**/

/**
	@var GenBank::Division GenBank::Synthetic
	The SYN division of GenBank containing Synthetic sequence data.
**/

/**
	@var GenBank::Division GenBank::Unannotated
	The UNA division of GenBank containing Unannotated sequence data.
**/

/**
	@var GenBank::Division GenBank::Patent
	The PAT division of GenBank containing Patent sequence data.
**/

/**
	@var GenBank::Division GenBank::Environmental
	The ENV division of GenBank containing Environmental sequence data.
**/

/**
	@var GenBank::Division GenBank::Constructed
	The CON division of GenBank containing Constructed sequence data.
**/

/**
	@var GenBank::Division GenBank::EST
	The EST division of GenBank containing Expressed sequence tag data.
**/

/**
	@var GenBank::Division GenBank::STS
	The STS division of GenBank containing Sequence tagged sites data.
**/

/**
	@var GenBank::Division GenBank::GSS
	The GSS division of GenBank containing Genome survey sequence data.
**/

/**
	@var GenBank::Division GenBank::HTGS
	The HTG division of GenBank containing High-throughput genomic sequence data.
**/

/**
	@var GenBank::Division GenBank::HTC
	The HTC division of GenBank containing High-throughput cDNA sequence data.
**/

/**
	@var GenBank::Division GenBank::NumDivisions
	The number of GenBank divisions - not a valid division.
**/

/**
	Default Constructor
**/
GenBank::GenBank() : m_release( 0 ), m_cmdId( 0 ), m_connected( false ), m_hasError( false ), m_ftp( NULL )
{
}

/**
	Destructor
**/
GenBank::~GenBank()
{
#ifdef QFTP
	delete m_ftp;
#endif
	m_ftp = NULL;
}

/**
	Connect to the GenBank ftp and get file and release information. Depending on 
	the connection, this operation may take 10 sec to 1 min or more.
	@return	true on success, false if the ftp server could not be contacted.
**/
bool GenBank::connectToGenBank()
{
	if( m_ftp )
		return !m_hasError;

#ifdef QFTP
	m_ftp = new QFtp;

	connect( m_ftp, SIGNAL(commandFinished(int,bool)), this, SLOT(commandFinished(int,bool)) );
	connect( m_ftp, SIGNAL(listInfo(const QUrlInfo&)), this, SLOT(listInfo(const QUrlInfo&)) );

	if( !execute( m_ftp->connectToHost( c_ftpHost ) ) ){
		Core::error( "Error connecting to GenBank ftp: " + m_ftp->errorString() );
		return false;
	}

	if( !execute( m_ftp->login() ) ){
		Core::error( "Error logging in to GenBank ftp: " + m_ftp->errorString() );
		return false;
	}
	
	if( !execute( m_ftp->cd( c_ftpDir ) ) ){
		Core::error( "Error changing directory in GenBank ftp: " + m_ftp->errorString() );
		return false;
	}

	if( !execute( m_ftp->list() ) ){
		Core::error( "Error listing directory in GenBank ftp: " + m_ftp->errorString() );
		return false;
	}

	QBuffer	releaseNumber;

	releaseNumber.open( QIODevice::WriteOnly );

	if( !execute( m_ftp->get( c_releaseNumberFile, &releaseNumber ) ) ){
		Core::error( "Error getting release number from GenBank ftp: " + m_ftp->errorString() );
		return false;
	}

	releaseNumber.close();

	QByteArray	arr = releaseNumber.buffer().trimmed();

	m_release = arr.toInt();
#endif

	return !m_hasError;
}

/**
	Get the GenBank release number. This is only valid after a successful 
	call to connectToGenBank().
	@return	The current release number for the GenBank flat files.
**/
int GenBank::getReleaseNumber() const
{
	return m_release;
}

/**
	Get the size (in bytes) of the given GenBank division. This function 
	calculates the total size for all the flat files in the given division.
	This is only valid after a successful call to connectToGenBank().
	@param	div	The GenBank division to calculate the size of.
	@return	The size (in bytes) of the GenBank division.
**/
qint64 GenBank::getDivisionSize( Division div ) const
{
	if( div >= Primate && div < NumDivisions )
		return m_data[div].m_totalSize;

	return 0;
}

/**
	Get the number of files for the given GenBank division. This function 
	gets the total number of flat files in the given division. This is only 
	valid after a successful call to connectToGenBank().
	@param	div	The GenBank division to get the number of files for.
	@return	The number of flat files in the GenBank division.
**/
int GenBank::getNumFiles( Division div ) const
{
	if( div >= Primate && div < NumDivisions )
		return m_data[div].m_files.count();

	return 0;
}

/**
	Get the name of a file in the given GenBank division. This function 
	returns the name of a flat file in the given division. Only the file 
	name is returned, relative to the GenBank ftp directory. This is only 
	valid after a successful call to connectToGenBank().
	@param	div		The GenBank division to get the name of a file for.
	@param	which	The index of the file to get the name of.
	@return	The name of a flat file in the GenBank division.
**/
QString GenBank::getFileName( Division div, int which ) const
{
	if( div < Primate || div >= NumDivisions )
		return QString::null;

	if( which < 0 || which >= m_data[div].m_files.count() )
		return QString::null;

	return m_data[div].m_files[which];
}

/**
	Get the URL of a file in the given GenBank division. This function 
	returns the URL of a flat file in the given division. The full URL 
	of the file is returned, including the 'ftp://' identifier. This is 
	only valid after a successful call to connectToGenBank().
	@param	div		The GenBank division to get the URL of a file for.
	@param	which	The index of the file to get the URL of.
	@return	The URL of a flat file in the GenBank division.
**/
QUrl GenBank::getFileUrl( Division div, int which ) const
{
	if( div < Primate || div >= NumDivisions )
		return QUrl();

	if( which < 0 || which >= m_data[div].m_files.count() )
		return QUrl();

	return QUrl( QString( "ftp://%1/%2/%3" ).arg( c_ftpHost ).arg( c_ftpDir ).arg( m_data[div].m_files[which] ) );
}

/**
	Get the size (in bytes) of a file in the given GenBank division. This 
	is only valid after a successful call to connectToGenBank().
	@param	div		The GenBank division to get the size of a file for.
	@param	which	The index of the file to get the size of.
	@return	The size (in bytes) of a flat file in the GenBank division.
**/
qint64 GenBank::getFileSize( Division div, int which ) const
{
	if( div < Primate || div >= NumDivisions )
		return 0;

	if( which < 0 || which >= m_data[div].m_fileSizes.count() )
		return 0;

	return m_data[div].m_fileSizes[which];
}

/**
	Get a user-readable short description of a GenBank division.
	@param	div		The GenBank division to get the description of.
	@return	A brief string describing the given division.
**/
QString GenBank::getDivisionDescription( Division div )
{
	if( div >= Primate && div < NumDivisions )
		return c_defs[div].m_desc;

	return QString::null;
}

/**
	Get the three-letter identifier of a GenBank division.
	@param	div		The GenBank division to get the abbreviation for.
	@return	A three-letter string abbreviation for the given division.
**/
QString GenBank::getDivisionAbbreviation( Division div )
{
	if( div >= Primate && div < NumDivisions )
		return c_defs[div].m_abbr;

	return QString::null;
}

/**
	Called by the ftp to notify us that the given command has finished.
	@param	id		The id of the command that finished.
	@param	error	true if an error occurred, false if the command executed successfully.
**/
void GenBank::commandFinished( int id, bool error )
{
	if( m_cmdId == id ){
		m_busy = false;
		m_hasError = error;
	}
	else {
		Core::warning( "Mismatched command Id in GenBank::commandFinished()" );
	}
}

/**
	Callback for the ftp class to pass in the information for a file on the ftp.
	@param	info	The details for a file on the ftp.
**/
void GenBank::listInfo( const QUrlInfo &info )
{
#ifdef QFTP
	if( !info.isFile() )
		return;

	QString	name = info.name();

	if( !name.startsWith( "gb" ) && name.endsWith( ".seq.gz" ) )
		return;

	QString	prefix = name.left( 5 );

	for( int i = 0; i < NumDivisions; i++ ){
		if( c_defs[i].m_filePrefix == prefix ){
			qint64	size = info.size();

			m_data[i].m_totalSize += size;
			m_data[i].m_files.append( name );
			m_data[i].m_fileSizes.append( size );
			break;
		}
	}
#endif
}

/**
	Called to wait for a command to synchronously execute on the ftp.
	@param	commandId	The command id returned by the the ftp function.
	@return	true if the command executed successfully, otherwise false.
**/
bool GenBank::execute( int commandId )
{
	m_busy = true;
	m_cmdId = commandId;

	while( m_busy )
		qApp->processEvents();

	return !m_hasError;
}

/**
	Default constructor
	Initializes the GenBank data class.
**/
GenBank::DivisionData::DivisionData() : m_totalSize( 0 )
{
}
